from mootdx.quotes import Quotes
import pandas as pd

# 获取所有股票列表
print("正在获取所有股票列表...")
client = Quotes.factory(market='std')
# 获取上海和深圳市场的股票列表
sh_stocks = client.stocks(market=1)  # 上海市场
sz_stocks = client.stocks(market=0)  # 深圳市场

print(f"上海市场股票数量: {len(sh_stocks)}")
print(f"深圳市场股票数量: {len(sz_stocks)}")

print("\n上海市场前10只股票:")
print(sh_stocks.head(10))

print("\n深圳市场前10只股票:")
print(sz_stocks.head(10))

# 合并并检查数据
all_stocks = pd.concat([sh_stocks, sz_stocks], ignore_index=True)
print(f"\n合并后总股票数量: {len(all_stocks)}")

# 查看是否有重复的代码
duplicated = all_stocks[all_stocks.duplicated('code', keep=False)]
print(f"\n重复的股票代码数量: {len(duplicated)}")

# 检查代码类型分布
all_stocks['code_prefix'] = all_stocks['code'].str[:3]
print("\n股票代码前3位分布:")
print(all_stocks['code_prefix'].value_counts().head(20))

# 筛选个股（只保留6位数字代码的个股）
# 普通个股代码规则：
# 沪市主板：600、601、603、605开头
# 深市主板：000开头
# 深市中小板：002开头
# 深市创业板：300开头（但根据要求需要过滤掉）
# 沪市科创板：688开头（但根据要求需要过滤掉）
# 北交所：8开头（但根据要求需要过滤掉）

# 先筛选出6位数字代码的证券
individual_stocks = all_stocks[all_stocks['code'].str.match(r'^\d{6}$')].copy()
print(f"\n6位数字代码证券数量: {len(individual_stocks)}")

# 再筛选出符合个股代码规则的证券（排除科创板、创业板和北交所）
# 个股代码应该以以下数字开头：
# 沪市主板: 600, 601, 603, 605
# 深市主板: 000
# 深市中小板: 002
filtered_stocks = individual_stocks[
    individual_stocks['code'].str.startswith(('600', '601', '603', '605', '000', '002'))
]
print(f"符合条件的个股数量: {len(filtered_stocks)}")

# 去除重复代码
filtered_stocks = filtered_stocks.drop_duplicates(subset=['code'])
print(f"去重后个股数量: {len(filtered_stocks)}")

# 显示筛选后的个股代码示例
print("\n符合条件的个股代码示例:")
print(filtered_stocks[['code', 'name']].head(20))

# 检查各个类别数量
sh_main_board = filtered_stocks[filtered_stocks['code'].str.startswith(('600', '601', '603', '605'))]
sz_main_board = filtered_stocks[filtered_stocks['code'].str.startswith('000')]
sz_sme_board = filtered_stocks[filtered_stocks['code'].str.startswith('002')]

print(f"\n沪市主板个股数量: {len(sh_main_board)}")
print(f"深市主板个股数量: {len(sz_main_board)}")
print(f"深市中小板个股数量: {len(sz_sme_board)}")
print(f"总计个股数量: {len(sh_main_board) + len(sz_main_board) + len(sz_sme_board)}")